import argparse
import os
import tarfile
import random
import shutil


def parse_args():
    parser = argparse.ArgumentParser(description='Sysc90K To 10K')
    parser.add_argument('tar_path', help='tar path')
    parser.add_argument('label_list_file', help='label list file')
    parser.add_argument('-o', '--out-dir', help='output path')
    parser.add_argument('--prefix', default='mnt/ramdisk/max/90kDICT32px', type=str, help='prefix path')
    # mnt/ramdisk/max/90kDICT32px
    parser.add_argument('--count', default=1000000, type=int, help='output image count')
    parser.add_argument('--has-extract', action='store_true', help='has extract')
    # extract
    args = parser.parse_args()
    return args


def main():
    args = parse_args()
    if not os.path.exists(args.out_dir):
        os.makedirs(args.out_dir)

    out_dir = args.out_dir
    prefix = args.prefix
    # with tarfile.open(args.tar_path) as file:
    #     file_name = '1/1_pontifically_58805.jpg'
    #     file_path = os.path.normpath(os.path.join(prefix, file_name))
    #     file.extract(file_path, path=out_dir)
    #     # file.list()
    #     # mnt/ramdisk/max/90kDICT32px/annotation_train.txt
    #     # mnt/ramdisk/max/90kDICT32px/annotation_val.txt
    #     # mnt/ramdisk/max/90kDICT32px/lexicon.txt
    #     # mnt/ramdisk/max/90kDICT32px/annotation_test.txt
    #     # mnt/ramdisk/max/90kDICT32px/annotation.txt
    #     # mnt/ramdisk/max/90kDICT32px/imlist.txt
    #     # mnt/ramdisk/max/90kDICT32px/3000/1/292_PRIMED_60081.jpg
    # return
    random.seed(123456)

    num_list = []
    no_num_list = []
    with open(args.label_list_file, 'r', encoding='utf-8') as file:
        # ./1/1/176_fieldwork_28792.jpg fieldwork
        for line in file.readlines():
            lines = line.strip().split()
            file_name = lines[0]
            label = lines[1]
            if label.isalnum() and (not label.isalpha()):  # 包含数字
                num_list.append([line, file_name])
            else:
                no_num_list.append([line, file_name])
    count = args.count
    num_count = len(num_list)
    half_count = count // 2
    if num_count > half_count:
        select_num_count = half_count
    else:
        select_num_count = num_count
    no_num_count = len(no_num_list)
    count = count - num_count
    if no_num_count >= count:
        select_no_num_count = count
    else:
        select_no_num_count = no_num_count
        if select_num_count < num_count:
            count = count - select_no_num_count
            left_num_count = num_count - select_num_count
            if left_num_count >= count:
                select_num_count += count
            else:
                select_num_count += left_num_count
    select_list = []
    if num_count > 0:
        random.shuffle(num_list)
        select_list.extend(num_list[:select_num_count])
        print('select_num_count:', select_num_count)
    if no_num_count > 0:
        random.shuffle(no_num_list)
        select_list.extend(no_num_list[:select_no_num_count])
        print('select_no_num_count:', select_no_num_count)

    if args.has_extract:
        out_select_list = []
        out_file_path = os.path.join(out_dir, 'out_label.txt')
        out_file = open(out_file_path, 'w', encoding='utf-8')
        file_count = 0
        src_dir = os.path.normpath(os.path.join(args.tar_path, prefix))
        dst_dir = os.path.normpath(os.path.join(out_dir, prefix))
        for line, file_name in select_list:
            src_file_path = os.path.join(src_dir, file_name)
            if not os.path.exists(src_file_path):
                print(src_file_path, 'Not Found.')
                continue
            dst_file_path = os.path.join(dst_dir, file_name)
            dir_name = os.path.dirname(dst_file_path)
            if not os.path.exists(dir_name):
                os.makedirs(dir_name)
            shutil.copy(src_file_path, dst_file_path)
            out_select_list.append(line)
            out_file.write(line)
            file_count += 1
            if file_count % 1000 == 0:
                print('file_count:', file_count)
                out_file.flush()
        if file_count % 1000 != 0:
            print('file_count:', file_count)
        out_file.close()
        out_select_list.sort()
        out_file_path = os.path.join(out_dir, 'label.txt')
        with open(out_file_path, 'w', encoding='utf-8') as file:
            for line in out_select_list:
                file.write(line)
    else:
        with tarfile.open(args.tar_path) as file:
            print('getmembers')
            members = file.getmembers()
            print('select_dict')
            select_dict = {}
            for line, file_name in select_list:
                file_path = os.path.normpath(os.path.join(prefix, file_name))
                select_dict[file_path] = line
            print('select_members')
            out_select_members_list = []
            for member in members:
                line = select_dict.pop(member.path, None)
                if line is not None:
                    out_select_members_list.append([line, member])
            print('out_select_list')
            out_select_members_list.sort(key=lambda x: x[0])
            select_members = []
            out_file_path = os.path.join(out_dir, 'label.txt')
            with open(out_file_path, 'w', encoding='utf-8') as out_file:
                for line, member in out_select_members_list:
                    out_file.write(line)
                    select_members.append(member)
            print('extract...')
            file.extractall(out_dir, members=select_members)
    print('Finish!')


if __name__ == '__main__':
    main()
